import torch.nn as nn
import torch.optim as optim

from config_manger.train_config import TrainConfig
from data_manager.loaders.data_loader import get_loader
from models_manger.le_net_model import LeNetModel
from trainer_manger.trainer_manger import TrainerManger


def test():
    train_loader, test_loader = get_loader("mnist")
    model = LeNetModel()
    print(model)
    epoch = TrainConfig.epochs
    criterion = nn.CrossEntropyLoss()  # 交叉熵损失
    optimizer = optim.Adam(model.parameters(), lr=TrainConfig.lr)  # Adam优化器
    trainer_manger = TrainerManger(model, optimizer, criterion, train_loader, test_loader, epoch)
    trainer_manger.train()


if __name__ == '__main__':
    test()
